/* Copyright 2013, Timothy Demulder <timothy@syphzero.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include <arpa/inet.h>
#include <ctype.h>
#include <fcntl.h>
#include <netinet/ether.h>
#include <netinet/if_ether.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <net/if_arp.h>
#include <pcap/pcap.h>
#include <regex.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>
#include "wurmd.h"

static char *bspth;

void sichk() {
	char *lckpth;
	int fd;
	struct flock fdfl;
	fdfl.l_type = F_WRLCK;
	fdfl.l_whence = SEEK_SET;
	fdfl.l_start=0;
	fdfl.l_len=0;
	bspth = getenv("HOME");
	if (bspth == NULL || bspth[0] != '/')
		eprintf("No home directory (or a bad one) defined.\n");
	lckpth = smalloc((sizeof(char)*strlen(bspth)) + (sizeof(char)*sizeof("/"LCKFD)));
	strcpy(lckpth,bspth);
	strcat(lckpth,"/"LCKFD);
	fd = open(lckpth,O_RDWR|O_CREAT,0600);
	if (fd == -1)
		eprintf("Can't create a lockfile.\n");
	if (fcntl(fd,F_SETLK,&fdfl) == -1)
		eprintf("Another instance is running.\n");
}

unsigned char *mkwolpkt(char *aeaddr) {
	int i,offset;
	unsigned char *pld=smalloc(sizeof(char)*102); //6+(16*6) bytes
	struct ether_addr *eaddr;
	eaddr = ether_aton(aeaddr);
	offset=0;
	memset(pld+offset, 0xff, sizeof(char)*6);
	offset=(sizeof(char))*6;
	for (i = 0; i < 16; i++) {
		memcpy(pld+offset,eaddr,sizeof(char)*6);
		offset+=6;
	}
	return pld;
}

int sendpkt(unsigned char *pkt) {
	const char *bcastinet="255.255.255.255";
	const int portnr=9;
	const int pktsize=(sizeof(char)*102);
	int sockptr = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sockptr < 0)
		eprintf("Couldn't create socket\n");
	int enbcast = 1;
	int ret = setsockopt(sockptr,SOL_SOCKET,SO_BROADCAST, &enbcast, sizeof(enbcast));
	if (ret) {
		close(sockptr);
		eprintf("Couldn't set broadcast flag on socket.");
	}
	struct sockaddr_in bcastaddr;
	memset(&bcastaddr, 0, sizeof(bcastaddr));
	bcastaddr.sin_family = AF_INET;
	inet_pton(AF_INET, bcastinet , &bcastaddr.sin_addr);
	bcastaddr.sin_port = htons(portnr);
	ret = sendto(sockptr,pkt,pktsize,0,(struct sockaddr*)&bcastaddr, sizeof bcastaddr);
	if (ret < 0 ) {
		close(sockptr); 
		return 1;
	}
	close(sockptr);
	return 0;
}

char *dtohaddr(const char *taddr) {
 	char *htaddr = smalloc(sizeof(char)*11);
	unsigned int htaddra, htaddrb, htaddrc, htaddrd;
	if (sscanf(taddr, "%d.%d.%d.%d", &htaddra, &htaddrb, &htaddrc, &htaddrd) != 4)
		eprintf("Couldn't convert address %s.\n",taddr);
	snprintf(htaddr,11,"0x%.2x%.2x%.2x%.2x", htaddra, htaddrb, htaddrc, htaddrd);
	return htaddr;
}

char *pktgettype(u_char *args,const struct pcap_pkthdr* pkthdr,const u_char* packet) {
	char *nmptr;
	const struct ether_header *etherhdr;
	etherhdr = (struct ether_header*)(packet);
	const struct ip *ip;
	ip = (struct ip*)(packet + sizeof(struct ether_header));
	const struct udphdr *udphdr;
	udphdr = (struct udphdr*)(packet + sizeof(struct ether_header) + sizeof(struct ip));
	const nbnshdr_t *nbnshdr;
	nbnshdr = (struct nbnshdr*)(packet + sizeof(struct ether_header) + sizeof(struct ip) + sizeof(struct udphdr));
	
	if (ntohs(etherhdr->ether_type) == ETHERTYPE_ARP) {
		struct arphdr *arp = (struct arphdr*)(packet + sizeof(struct ether_header));
		if(arp->ar_op == 256) {
			u_int32_t addr;	
			memcpy(&addr,ar_tpa(arp),sizeof(addr));
			nmptr=intoa(addr);
			vprintf("ARP request for %s\n",nmptr);
		}
	}

	/* nbnshdr->opcode should be 0x0001000000000001 for a NB name request */
	else if (ip->ip_p == IPPROTO_UDP && ntohs(udphdr->dest) == 137 && nbnshdr->opcode == 4097) {
		unsigned char databuf,datanbuf;
		char nbnsaddrbuf[16]="";
		int idx=0;
		const unsigned char *nbnsdata = (packet + sizeof(struct ether_header) + sizeof(struct ip) + sizeof(struct udphdr) + sizeof(struct nbnshdr) -1);
		for(;;) {
			databuf = *nbnsdata;
			if (databuf=='\0' || databuf < 'A' || databuf > 'Z' )
				break;
			databuf-='A';
			datanbuf=databuf<<4;
			nbnsdata++;
			databuf = *nbnsdata;
			if (databuf=='\0' || databuf < 'A' || databuf > 'Z' )
				eprintf("Malformed NBNS packet, something is very wrong here.\n");
			databuf-='A';
			datanbuf= databuf | datanbuf;
			nbnsdata++;
			if (datanbuf==32)
				continue;
			if (idx <= NBNS_NAME_MAX)
				nbnsaddrbuf[idx++]=datanbuf;
		}
		nbnsaddrbuf[idx]='\0';
		nmptr=nbnsaddrbuf;
		vprintf("NBNS request for %s detected\n",nmptr);
	}

	else {
		char taddrbuf[INET_ADDRSTRLEN];
		inet_ntop(AF_INET,&(ip->ip_dst),taddrbuf,INET_ADDRSTRLEN);
		nmptr = taddrbuf;
		vprintf("Initial TCP connection detected to %s\n",nmptr);
	}
	return nmptr;
	
}

char * intoa(u_int32_t addr) { // Stolen from tcpdump
        register char *cp;
        register u_int byte;
        register int n;
        static char buf[sizeof(".xxx.xxx.xxx.xxx")];
        NTOHL(addr);
        cp = buf + sizeof(buf);
        *--cp = '\0';
        n = 4;
        do {
                byte = addr & 0xff;
                *--cp = byte % 10 + '0';
                byte /= 10;
                if (byte > 0) {
                        *--cp = byte % 10 + '0';
                        byte /= 10;
                        if (byte > 0)
                                *--cp = byte + '0';
                }
                *--cp = '.';
                addr >>= 8;
        } while (--n > 0);
        return cp + 1;
}

void sndwolpktot(char *eaddr) {
	unsigned char *pkt=mkwolpkt(eaddr);
	if(sendpkt(pkt))
		eprintf("Couldn't broadcast WoL for %s\n",eaddr);
	vprintf("WoL broadcast for %s\n",eaddr);
}

char *pktgeteaddr(char *inpt) {
	if (inpt == '\0' || inpt == NULL)
		eprintf("Comparing something to nothing is no better than dividing by 0!\n");
	FILE *fp;
	char *bfr,*taddr,*eaddr,*nbaddr;
	static int rdbfr = 96;
	bfr = smalloc(sizeof(char)*rdbfr);
	taddr = smalloc(sizeof(char)*16);
	eaddr = smalloc(sizeof(char)*18);
	nbaddr = smalloc(sizeof(char)*17);
	fp = fopen(cnfile, "r");
	if(fp==NULL)
		eprintf("Can't open configuration file %s.\n",cnfile);
	while(fgets(bfr, rdbfr, fp)!=NULL) {
		if (!strncmp(bfr,"#",1) || !strncmp(bfr,"\n",1))
			continue;
		sscanf(bfr,"%s %s %s", taddr,eaddr,nbaddr);
		if (!strcmp(taddr,inpt))
			break;
		if (!strcmp(nbaddr,inpt))
			break;
	}
	fclose(fp);
	free(bfr);
	free(taddr);
	free(nbaddr);
	return(eaddr);
}

void pktmatch(u_char *args,const struct pcap_pkthdr* pkthdr,const u_char* packet) {
	char *nmptr=pktgettype(args,pkthdr,packet);
	char *eaddr=pktgeteaddr(nmptr);
	eaddrval(eaddr);
	unsigned char *wolpkt;
	wolpkt=mkwolpkt(eaddr);
	if (sendpkt(wolpkt))
		vprintf("couldn't broadcast WoL for %s.",eaddr);
	free(wolpkt);
	vprintf("WoL broadcast for %s.\n",eaddr);
	free(eaddr);
}

void eaddrval(char *eaddr) {
	regex_t re;
	if(regcomp(&re,"^\\([0-9a-fA-F]\\{2\\}:\\)\\{5\\}[0-9a-fA-F]\\{2\\}$",0))
		eprintf("Could not compile regex\n");
	if(regexec(&re,eaddr,0, NULL,0))
		eprintf("%s is not a valid mac address\n",eaddr);
	regfree(&re);
}

void taddrval(char *taddr) {
	struct in_addr tcstr;
	if (!inet_aton(taddr, &tcstr))
		eprintf("%s is not a valid IP address\n",taddr);
}

char *srealloc(char *optr, char *taptr) {
	optr = realloc(optr,strlen(optr)+(sizeof(char)*strlen(taptr)+1));
	if ( optr == NULL )
		eprintf("Couldn't assign needed memory.\n");
	return optr;
}

void *smalloc(size_t size) {
	void *ptr = malloc(size);
	if (ptr == NULL)
		eprintf("No memory available.\n");
	return ptr;
}

void sighndl(int sig) {
	signal(sig, SIG_IGN);
	vprintf("\nSIGINT detected, shutting down.\n");
	exit(0);
}
